#include "distribute_download.h"
#include <custom_head.h>

typedef struct {
  char ip[16];
  int port;
  int chunk_index;
} ServerInfo;

ServerInfo servers[MAX_SERVERS];
int total_chunks;
char output_file[256];
char temp_dir[] = "./tmp_downloads";

void *download_thread(void *arg) {
  ServerInfo *info = (ServerInfo *)arg;

  // 创建临时文件名
  char temp_file[256];
  sprintf(temp_file, "%s/chunk_%d.tmp", temp_dir, info->chunk_index);

  int sock = socket(AF_INET, SOCK_STREAM, 0);
  struct sockaddr_in addr;
  memset(&addr, 0, sizeof(addr));
  addr.sin_family = AF_INET;
  addr.sin_port = htons(info->port);
  inet_pton(AF_INET, info->ip, &addr.sin_addr);

  
  if (connect(sock, (struct sockaddr *)&addr, sizeof(addr))) {
    fprintf(stderr, "connect error\n");
    return NULL;
  }

  // 发送分块请求
  char request[32];
  sprintf(request, "GET %d", info->chunk_index);
  send(sock, request, strlen(request), 0);

  // 接收数据并写入临时文件
  int fd = open(temp_file, O_RDWR);
  if (!fd) {
    fprintf(stderr, "open file error.\n");
    close(sock);
    return NULL;
  }

  int total = 0;
  char buffer[CHUNK_SIZE];
  while (1) {
    int bytes = recv(sock, buffer, CHUNK_SIZE, 0);
    if (bytes <= 0)
      break;
    write(fd,buffer, bytes);
    total += bytes;
  }

  close(fd);
  close(sock);
  printf("client blocks %d has downloaded (%d Bytes)\n", info->chunk_index, total);
  return NULL;
}

// 合并所有分块
void merge_chunks() {
  int output = open(output_file,O_RDWR);
  if (!output) {
    fprintf(stderr, "open file err.\n");
    return;
  }

  for (int i = 0; i < total_chunks; i++) {
    char temp_file[256];
    sprintf(temp_file, "%s/chunk_%d.tmp", temp_dir, i);

    int chunk = open(temp_file, O_RDWR);
    if (!chunk) {
      printf("client blocks %d miss\n", i);
      continue;
    }

    // 写入最终文件
    char buffer[CHUNK_SIZE];
    int bytes;
    while ((bytes = read(chunk, buffer, CHUNK_SIZE)) > 0) {
      write(output ,buffer, bytes);
    }
    close(chunk);
    remove(temp_file); // 删除临时文件
  }

  close(output);
  printf("client all blocks has combine to %s\n", output_file);
}

int main(int argc, char *argv[]) {
  if (argc < 4) {
    printf("用法: %s <输出文件> <总分块数> <服务器1 IP:端口:分块号> ...\n",
           argv[0]);
    return 1;
  }

  // 初始化参数
  strcpy(output_file, argv[1]);
  total_chunks = atoi(argv[2]);
  int num_servers = argc - 3;
  mkdir(temp_dir, 0755); // 创建临时目录

  // 解析服务器信息
  for (int i = 0; i < num_servers; i++) {
    char *token = strtok(argv[i + 3], ":");
    strcpy(servers[i].ip, token);

    token = strtok(NULL, ":");
    servers[i].port = atoi(token);

    token = strtok(NULL, ":");
    servers[i].chunk_index = atoi(token);
  }

  // 创建下载线程
  pthread_t threads[MAX_SERVERS];
  for (int i = 0; i < num_servers; i++) {
    pthread_create(&threads[i], NULL, download_thread, &servers[i]);
  }

  // 等待所有线程完成
  for (int i = 0; i < num_servers; i++) {
    pthread_join(threads[i], NULL);
  }

  // 合并文件
  merge_chunks();
  rmdir(temp_dir); // 删除临时目录
  return 0;
}